from qt import *
from hexedit import HexViewDataSource, HexEdit
from tablewindow import TableWindow
from ollydbg import disasm

class MemoryDataSource(HexViewDataSource):

    def __init__(self, process):
        self.process = process

    def read(self, size, pos):
        try:
            return self.process.read(pos, size)
        except:
            return '\x00'*size

    def write(self, data, pos):
        self.process.write(pos, data)

    def length(self):
        return (1<<32)


            
class DisasmWindow(TableWindow):

    def __init__(self, parent=None):
        super(DisasmWindow, self).__init__(parent, ['Address', 'Hex', 'Disassembly', 'Comment'])


class RegWindow(TableWindow):

    def __init__(self, parent=None):
        super(RegWindow, self).__init__(parent, ['Name', 'Value', 'Comment'])


class InfoWindow(TableWindow):

    def __init__(self, parent=None):
        super(InfoWindow, self).__init__(parent, ['Info'], rows=2)
        self.table.horizontalHeader().setVisible(False)

        
class DumpWindow(TableWindow):

    def __init__(self, parent=None):
        super(DumpWindow, self).__init__(parent, ['Address', 'Dump', 'Str'])
        
        self.setCentralWidget(HexEdit())

class StackWindow(DumpWindow):
    pass

class CpuWindow(QMainWindow):

    def __init__(self, parent=None):
        super(CpuWindow, self).__init__(parent)
        self.setWindowTitle('CPU')

##               |
##        disasm | reg
##               1
##        ---0---|
##        info   |
##        -------3------
##               |
##        dump   2 stack
##               |

        spl0 = QSplitter(self)
        disasm = self.disasm = DisasmWindow(self)
        spl0.addWidget(disasm)
        spl0.addWidget(InfoWindow(self))
        spl0.setStretchFactor(0, 8)
        spl0.setStretchFactor(1, 2)
        spl0.setOrientation(Qt.Vertical)

        spl1 = QSplitter(self)
        spl1.addWidget(spl0)
        reg = self.reg = RegWindow(self)
        spl1.addWidget(reg)

        spl2 = QSplitter(self)
        dump = self.dump = DumpWindow(self)
        stack = self.stack = StackWindow(self)
        spl2.addWidget(dump)
        spl2.addWidget(stack)

        spl3 = QSplitter(self)
        spl3.addWidget(spl1)
        spl3.addWidget(spl2)
        spl3.setOrientation(Qt.Vertical)

        spl0.setFrameStyle(QFrame.Panel)
        spl1.setFrameStyle(QFrame.Panel)
        spl2.setFrameStyle(QFrame.Panel)
        spl3.setFrameStyle(QFrame.Panel)

        self.setCentralWidget(spl3)

    def show_dump(self, process, pc):
        hexedit = self.dump.centralWidget()
        hexedit.set_data_source(MemoryDataSource(process))
        hexedit.hex_view.update_data_and_cursor(pc)

    def show_disasm(self, process, thread, pc):
        self.disasm.delete_all_items()
        
        offset = 0
        while offset < 256:
            addr, size, disasm, dump  = thread.disassemble( pc+offset, 0x10 ) [0]
            addr = '%08X' % addr
            offset += size
            self.disasm.add_item(addr, dump, disasm)

    def show_stack(self, process, sp):
        hexedit = self.stack.centralWidget()
        hexedit.set_data_source(MemoryDataSource(process))
        hexedit.hex_view.update_data_and_cursor(sp)

    def show_reg(self, ctx):
        self.reg.delete_all_items() 
        for key in 'Eax Ecx Edx Ebx Esp Ebp Esi Edi'.split():
            value = ctx.get(key)
            self.reg.add_item(key, '%08X'%value)

            
# TODO: compile distrom3 and remove this hack

def Decode(address, code, mode):
    text, size = disasm(code, address)
    return [(address, size, text, code[:size].encode('hex').upper())]

from winappdbg import system
system.Decode = Decode

